// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Author: wsfuyibing <682805@qq.com>
// Date: 2024-11-20

package elastic

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"github.com/elastic/go-elasticsearch/v7"
	"github.com/elastic/go-elasticsearch/v7/esapi"
	"io"
	"viewer/app/search/searched"
)

// Searcher
// 查询工具.
type (
	Searcher interface {
		Mapper() searched.Mapper
		Search(ctx context.Context, tid, sid, pid string) (err error)
	}

	searcher struct {
		mapper searched.Mapper
	}
)

// New
// 创建查询工具.
func New() Searcher {
	return &searcher{
		mapper: searched.NewMapper(),
	}
}

func (o *searcher) Mapper() searched.Mapper {
	return o.mapper
}

// Search
// 查询数据.
func (o *searcher) Search(ctx context.Context, tid, sid, pid string) (err error) {
	var (
		body, buf []byte
		cli       *elasticsearch.Client
		item      *searched.Item
		req       = esapi.SearchRequest{}
		res       *esapi.Response
		yes       = true
	)

	// 1. 获取连接.
	if cli, err = conf.Client(); err != nil {
		return
	}

	// 2. 查询条件.
	must := make([]map[string]map[string]any, 0)

	// 2.1 链路ID.
	if tid != "" {
		must = append(must, map[string]map[string]any{"term": {"trace_id": tid}})
	}

	if sid != "" {
		must = append(must, map[string]map[string]any{"term": {"span_id": sid}})
	}

	if pid != "" {
		must = append(must, map[string]map[string]any{"term": {"parent_span_id": pid}})
	}

	// 2.2 附加参数.
	param := map[string]interface{}{
		"query": map[string]map[string]interface{}{"bool": {"must": must}},
		"from":  0,
		"size":  1000,
		"sort":  []string{},
		"aggs":  map[string]interface{}{},
	}

	// 3. 入参编码.
	if buf, err = json.Marshal(param); err == nil {
		req.Index = conf.Indexes
		req.Body = bytes.NewBuffer(buf)
		req.DocumentType = conf.IndexTypes
		req.AllowNoIndices = &yes
		req.AllowPartialSearchResults = &yes
		req.IgnoreUnavailable = &yes

		if res, err = req.Do(ctx, cli); err == nil {
			body, err = io.ReadAll(res.Body)
		}
	}
	if err != nil {
		return
	}

	// 4. 解析结果.
	sr := &Result{}
	if err = json.Unmarshal(body, sr); err != nil {
		return
	}

	// 4.1 无效结果.
	if sr.Hits == nil {
		err = fmt.Errorf(`invalid elasticsearch result`)
		return
	}

	// 4.2 遍历结果.
	for _, hit := range sr.Hits.Hits {
		if buf, err = json.Marshal(hit.Source); err != nil {
			return
		}

		if item, err = searched.NewItem(buf); err != nil {
			return
		}

		o.mapper.Add(item)
	}
	return
}
